#2.13
import numpy as np
import matplotlib.pyplot as plt
# 模型参数
c = 1.0  # 波的传播速率
tc, te = 0.0, 1.0  # 时间范围，0<t<te
xa, xb = 0.0, 1.0  # 空间范围，xa<x<xb
ya, yb = 0.0, 1.0  # 空间范围，ya<y<yb
# 初始化
c2 = c*c  # 方程参数
dt = 0.01  # 时间步长
dx = dy = 0.02  # 空间步长
tNodes = round(te/dt)  # t轴 时序网格数
xNodes = round((xb-xa)/dx)  # $\mathrm{x}$轴 空间网格数
yNodes = round((yb-ya)/dy)  # $\mathrm{y}$轴 空间网格数
tZone = np.arange(0, (tNodes+1)*dt, dt)  # 建立空间网格
xZone = np.arange(0, (xNodes+1)*dx, dx)  # 建立空间网格
yZone = np.arange(0, (yNodes+1)*dy, dy)  # 建立空间网格
xx, yy = np.meshgrid(xZone, yZone)  # 生成网格点的坐标 xx,yy (二维数组)
# 步长比检验(r>1 则算法不稳定)
r = 4 * c2 * dt*dt / (dx*dx+dy*dy)
print("dt = {:.2f}, dx = {:.2f}, dy = {:.2f}, r = {:.2f}".format(dt,dx,dy,r))
assert r < 1.0, "Error: r>1, unstable step ratio of dt2/(dx2+dy2) ."
rx = c*c * dt**2/dx**2
ry = c*c * dt**2/dy**2
# 绘图
fig = plt.figure(figsize=(8,6))
ax1 = fig.add_subplot(111, projection='3d')
# 计算初始值
U = np.zeros([tNodes+1, xNodes+1, yNodes+1])  # 建立三维数组
U[0] = np.sin(6*np.pi*xx)+np.cos(4*np.pi*yy)  # U[0,:,:]
U[1] = np.sin(6*np.pi*xx)+np.cos(4*np.pi*yy)  # U[1,:,:]
surf = ax1.plot_surface(xx, yy, U[0,:,:], rstride=2, cstride=2, cmap=plt.cm.coolwarm)
# 有限差分法求解
for k in range(2,tNodes+1):
    for i in range(1,xNodes):
        for j in range(1,yNodes):
            U[k,i,j] = rx*(U[k-1,i-1,j]+U[k-1,i+1,j]) + ry*(U[k-1,i,j-1]+U[k-1,i,j+1])\
                     + 2*(1-rx-ry)*U[k-1,i,j] -U[k-2,i,j]
surf = ax1.plot_surface(xx, yy, U[k,:,:], rstride=2, cstride=2, cmap='rainbow')
ax1.set_xlim3d(0, 1.0)
ax1.set_ylim3d(0, 1.0)
ax1.set_zlim3d(-2, 2)
ax1.set_title("2D wave equationt (t= %.2f)" % (k*dt))
ax1.set_xlabel("x")
ax1.set_ylabel("y")
plt.show()
